import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gamma


def matern_cluster_process(
    center, radius, num_clusters, lambda_p, mean_cluster_size, gamma_shape
):
    """
    模拟Matern簇过程

    :param center: 圆形区域的中心坐标 (x0, y0)
    :param radius: 圆形区域的半径
    :param num_clusters: 母点（簇中心）的数量
    :param lambda_p: 母点过程的强度
    :param mean_cluster_size: 平均簇大小
    :param gamma_shape: 用于生成簇半径的伽马分布的形状参数
    :return: 模拟得到的所有点的坐标数组，簇中心点的x坐标数组，簇中心点的y坐标数组
    """
    # 生成母点（簇中心）的坐标
    parent_points_x = np.random.uniform(
        center[0] - radius, center[0] + radius, num_clusters
    )
    parent_points_y = np.random.uniform(
        center[1] - radius, center[1] + radius, num_clusters
    )

    all_points = []

    for i in range(num_clusters):
        # 生成簇半径
        cluster_radius = gamma.rvs(gamma_shape)

        # 确定簇内点的数量，假设服从泊松分布
        num_points_in_cluster = np.random.poisson(mean_cluster_size)

        # 生成簇内点的坐标
        points_in_cluster_x = np.random.uniform(
            parent_points_x[i] - cluster_radius,
            parent_points_x[i] + cluster_radius,
            num_points_in_cluster,
        )
        points_in_cluster_y = np.random.uniform(
            parent_points_y[i] - cluster_radius,
            parent_points_y[i] + cluster_radius,
            num_points_in_cluster,
        )

        # 筛选出在圆形区域内的点
        valid_indices = np.where(
            (points_in_cluster_x - center[0]) ** 2
            + (points_in_cluster_y - center[1]) ** 2
            <= radius**2
        )
        valid_points_x = points_in_cluster_x[valid_indices]
        valid_points_y = points_in_cluster_y[valid_indices]

        all_points.extend([(x, y) for x, y in zip(valid_points_x, valid_points_y)])

    return np.array(all_points), parent_points_x, parent_points_y


# 设置参数
center = (0, 0)  # 圆形区域中心坐标
radius = 20  # 圆形区域半径
num_clusters = 20  # 母点（簇中心）数量
lambda_p = 0.5  # 母点过程的强度
mean_cluster_size = 8  # 平均簇大小
gamma_shape = 2  # 用于生成簇半径的伽马分布的形状参数

# 模拟Matern簇过程
points, parent_points_x, parent_points_y = matern_cluster_process(
    center, radius, num_clusters, lambda_p, mean_cluster_size, gamma_shape
)

# 提取簇中心点坐标
parent_points = np.array(
    [(parent_points_x[i], parent_points_y[i]) for i in range(num_clusters)]
)

# 绘制图形
plt.scatter(points[:, 0], points[:, 1], s=10)
# 绘制簇中心点，用三角形表示
plt.scatter(parent_points[:, 0], parent_points[:, 1], marker="^", s=50, c="green")
plt.xlim(center[0] - radius, center[0] + radius)
plt.ylim(center[1] - radius, center[1] + radius)
plt.gca().set_aspect("equal", adjustable="box")
plt.title("Matern Cluster Process Simulation")
plt.xlabel("X-axis")
plt.ylabel("Y-axis")
plt.show()
